"""
LLM客户端模块
用于调用LLM API生成代码模块总结
"""
import requests
import json
from typing import List, Dict, Optional, Tuple
from config import LLMConfig
from concurrent.futures import ThreadPoolExecutor, as_completed
import time


class LLMClient:
    """LLM API客户端"""
    
    def __init__(self, config: LLMConfig):
        self.config = config
        self.session = requests.Session()
        self.session.headers.update({
            "Content-Type": "application/json",
            "Authorization": f"Bearer {config.api_key}"
        })
        # 并发控制
        self.max_concurrent_requests = getattr(config, 'max_concurrent_requests', 5)  # 最大并发请求数
        self.request_delay = getattr(config, 'request_delay', 0.1)  # 请求间隔（秒）
    
    def summarize_code_module(self, code_snippets: List[str], 
                            context: Optional[Dict] = None) -> str:
        """
        调用LLM总结代码模块功能
        
        Args:
            code_snippets: 代码片段列表
            context: 额外上下文信息
        
        Returns:
            总结文本
        """
        if not code_snippets:
            return "没有代码片段可供分析"
        
        # 构建prompt
        prompt = self._build_summary_prompt(code_snippets, context)
        
        # 调用API
        try:
            response = self._call_api(prompt)
            return self._parse_response(response)
        except requests.exceptions.Timeout:
            print("LLM API调用超时")
            return self._fallback_summary(code_snippets)
        except requests.exceptions.RequestException as e:
            print(f"LLM API调用失败: {e}")
            return self._fallback_summary(code_snippets)
        except Exception as e:
            print(f"处理LLM响应时出错: {e}")
            return self._fallback_summary(code_snippets)
    
    def summarize_code_modules_batch(self, patterns: List[Dict]) -> List[Tuple[int, str]]:
        """
        批量并发调用LLM总结多个代码模块
        
        Args:
            patterns: 模式列表，每个包含code_snippets和context
        
        Returns:
            [(index, summary), ...] 返回索引和对应的总结
        """
        if not patterns:
            return []
        
        print(f"开始批量生成{len(patterns)}个模式的总结...")
        start_time = time.time()
        
        results = []
        
        # 使用线程池并发处理
        with ThreadPoolExecutor(max_workers=self.max_concurrent_requests) as executor:
            # 提交所有任务
            future_to_index = {}
            
            for i, pattern in enumerate(patterns):
                # 为避免过载，添加小延迟
                if i > 0 and i % self.max_concurrent_requests == 0:
                    time.sleep(self.request_delay)
                
                # 提交任务
                future = executor.submit(
                    self._summarize_with_retry,
                    pattern.get('code_snippets', []),
                    pattern.get('context', {})
                )
                future_to_index[future] = i
            
            # 收集结果
            completed = 0
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    summary = future.result()
                    results.append((index, summary))
                    completed += 1
                    if completed % 10 == 0:
                        print(f"已完成 {completed}/{len(patterns)} 个总结")
                except Exception as e:
                    print(f"处理模式 {index} 时出错: {e}")
                    # 使用备用总结
                    results.append((index, self._fallback_summary(
                        patterns[index].get('code_snippets', [])
                    )))
        
        # 按原始顺序排序
        results.sort(key=lambda x: x[0])
        
        elapsed = time.time() - start_time
        print(f"批量生成完成，耗时 {elapsed:.1f} 秒")
        
        return results
    
    def _summarize_with_retry(self, code_snippets: List[str], 
                             context: Optional[Dict] = None, 
                             max_retries: int = 2) -> str:
        """
        带重试的总结方法
        """
        for attempt in range(max_retries + 1):
            try:
                prompt = self._build_summary_prompt(code_snippets, context)
                response = self._call_api(prompt)
                return self._parse_response(response)
            except requests.exceptions.Timeout:
                if attempt < max_retries:
                    time.sleep(1 * (attempt + 1))  # 递增延迟
                    continue
                return self._fallback_summary(code_snippets)
            except Exception as e:
                if attempt < max_retries:
                    time.sleep(0.5)
                    continue
                return self._fallback_summary(code_snippets)
    
    def _build_summary_prompt(self, code_snippets: List[str], 
                             context: Optional[Dict]) -> str:
        """构建总结prompt"""
        # 限制代码片段数量和长度
        max_snippets = 5
        max_length_per_snippet = 500
        
        snippets_text = []
        for i, code in enumerate(code_snippets[:max_snippets]):
            # 截断过长的代码
            if len(code) > max_length_per_snippet:
                code = code[:max_length_per_snippet] + "..."
            snippets_text.append(f"代码片段{i+1}:\n```python\n{code}\n```")
        
        code_text = "\n\n".join(snippets_text)
        
        prompt = f"""请分析以下Python测试代码片段，这些代码片段在多个测试文件中重复出现。
请总结它们的共同功能模式，注意不要涉及参数，因为参数在其他的模式中可能改变。尽量总结出通用的功能模式。

{code_text}

请用简洁的中文描述：
1. 这些代码片段的共同功能是什么
2. 结合使用的命令，给出操作步骤,并给出该命令，比如使用...命令...

输出格式：
功能：[一句话描述，不超过100字]
步骤：[用1-2-3...的格式列出步骤，每步不超过30字]

要求：
- 描述简洁明了，不要输出除功能和步骤外的其他内容
"""

        if context:
            # 添加统计信息
            if 'frequency' in context:
                prompt += f"\n\n补充信息：此模式在代码库中出现了{context['frequency']}次"
            if 'pattern_type' in context:
                prompt += f"，模式类型为{context['pattern_type']}"
        
        return prompt
    
    def _call_api(self, prompt: str) -> Dict:
        """调用LLM API"""
        payload = {
            "model": self.config.model,
            "messages": [
                {
                    "role": "system", 
                    "content": "你是一个专业的Python代码分析专家，擅长总结测试代码的功能模式。请用简洁准确的中文描述代码功能。"
                },
                {
                    "role": "user", 
                    "content": prompt
                }
            ],
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature,
            "top_p": self.config.top_p,
            "top_k": self.config.top_k,
        }
        
        response = self.session.post(
            self.config.api_url,
            json=payload,
            timeout=self.config.timeout
        )
        
        # 检查响应状态
        if response.status_code != 200:
            raise requests.exceptions.RequestException(
                f"API返回错误状态码: {response.status_code}, 响应: {response.text}"
            )
        
        return response.json()
    
    def _parse_response(self, response: Dict) -> str:
        """解析API响应"""
        try:
            # 尝试从响应中提取内容
            if 'choices' in response and len(response['choices']) > 0:
                message = response['choices'][0].get('message', {})
                content = message.get('content', '')
                if "</think>" in content:
                    content = content.split("</think>")[-1]
                if content:
                    return content.strip()
            
            # 如果响应格式不符合预期
            return "无法解析LLM响应"
            
        except (KeyError, IndexError, TypeError) as e:
            print(f"解析响应时出错: {e}")
            return "响应格式错误"
    
    def generate_test_cases_batch(self, scripts: List[Dict]) -> List[Tuple[int, Dict]]:
        """
        批量生成测试用例
        
        Args:
            scripts: 脚本列表，每个包含script_content和script_path
        
        Returns:
            [(index, test_case), ...] 返回索引和对应的测试用例
        """
        if not scripts:
            return []
        
        print(f"开始批量生成{len(scripts)}个脚本的测试用例...")
        start_time = time.time()
        
        results = []
        
        # 使用线程池并发处理
        with ThreadPoolExecutor(max_workers=self.max_concurrent_requests) as executor:
            # 提交所有任务
            future_to_index = {}
            
            for i, script in enumerate(scripts):
                # 为避免过载，添加小延迟
                if i > 0 and i % self.max_concurrent_requests == 0:
                    time.sleep(self.request_delay)
                
                # 提交任务
                future = executor.submit(
                    self._generate_test_case_with_retry,
                    script.get('script_content', ''),
                    script.get('script_path', '')
                )
                future_to_index[future] = i
            
            # 收集结果
            completed = 0
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    test_case = future.result()
                    results.append((index, test_case))
                    completed += 1
                    if completed % 10 == 0:
                        print(f"已完成 {completed}/{len(scripts)} 个测试用例")
                except Exception as e:
                    print(f"处理脚本 {index} 时出错: {e}")
                    # 使用备用测试用例
                    results.append((index, self._fallback_test_case(
                        scripts[index].get('script_content', ''),
                        scripts[index].get('script_path', '')
                    )))
        
        elapsed_time = time.time() - start_time
        print(f"批量生成完成，耗时: {elapsed_time:.2f}秒")
        
        # 按原始索引排序
        results.sort(key=lambda x: x[0])
        
        return results
    
    def _generate_test_case_with_retry(self, script_content: str, script_path: str, max_retries: int = 3) -> Dict:
        """
        带重试的测试用例生成
        """
        for attempt in range(max_retries):
            try:
                prompt = self._build_test_case_prompt(script_content, script_path)
                response = self._call_api(prompt)
                content = self._parse_response(response)
                
                # 解析测试用例
                test_case = self._parse_test_case_response(content)
                return test_case
                
            except Exception as e:
                if attempt < max_retries - 1:
                    time.sleep(self.request_delay * (attempt + 1))
                else:
                    raise e
        
        return self._fallback_test_case(script_content, script_path)
    
    def _build_test_case_prompt(self, script_content: str, script_path: str) -> str:
        """构建生成测试用例的prompt"""
        import os
        
        # 截断过长的代码
        max_length = 8000
        if len(script_content) > max_length:
            script_content = script_content[:max_length] + "\n... (代码过长，已截断) ..."
        
        prompt = f"""请分析以下Python测试脚本，并生成对应的测试用例文档。

脚本文件名：{os.path.basename(script_path)}

脚本内容：
```python
{script_content}
```

请生成测试用例，格式要求如下：
1. 简要描述：一句话说明测试目的（不超过50字）
2. 测试步骤和预期结果：列出每个测试步骤及其预期结果
3. 每个测试步骤需要结合使用的主要命令来给出，并给出该命令，比如使用...命令...

请以JSON格式输出，格式如下：
{{
    "description": "测试描述",
    "test_steps": [
        {{
            "step": "步骤1描述",
            "expected_result": "预期结果1"
        }},
        {{
            "step": "步骤2描述", 
            "expected_result": "预期结果2"
        }}
    ]
}}

要求：
- 从脚本的docstring、注释和代码逻辑中提取测试信息
- 步骤描述要清晰具体，包含具体的操作和配置
- 预期结果要明确，说明应该观察到什么现象
- 如果有多个测试方法，将它们的步骤都包含在内
- 返回的JSON格式要符合要求，不要输出除测试用例外的其他内容
"""
        
        return prompt
    
    def _parse_test_case_response(self, content: str) -> Dict:
        """解析LLM响应为结构化测试用例"""
        try:
            # 尝试直接解析JSON
            if '{' in content and '}' in content:
                # 提取JSON部分
                start_idx = content.find('{')
                end_idx = content.rfind('}') + 1
                json_str = content[start_idx:end_idx]
                test_case = json.loads(json_str)
                
                # 验证必要字段
                if 'description' in test_case and 'test_steps' in test_case:
                    return test_case
        except:
            pass
        
        # 如果JSON解析失败，尝试从文本中提取
        test_case = {
            "description": "",
            "test_steps": []
        }
        
        lines = content.split('\n')
        for line in lines:
            line = line.strip()
            if line.startswith('描述:') or line.startswith('功能:'):
                test_case['description'] = line.split(':', 1)[1].strip()
            elif line.startswith('步骤') and ':' in line:
                # 提取步骤
                step_text = line.split(':', 1)[1].strip()
                test_case['test_steps'].append({
                    "step": step_text,
                    "expected_result": "验证操作成功"
                })
        
        return test_case
    
    def _fallback_test_case(self, script_content: str, script_path: str) -> Dict:
        """降级方案：基于规则生成测试用例"""
        import os
        import re
        
        test_case = {
            "description": f"测试脚本 {os.path.basename(script_path)}",
            "test_steps": []
        }
        
        # 尝试从docstring提取描述
        docstring_match = re.search(r'"""(.*?)"""', script_content, re.DOTALL)
        if docstring_match:
            desc_lines = docstring_match.group(1).strip().split('\n')
            for line in desc_lines:
                if '测试目的' in line or '测试功能' in line:
                    test_case['description'] = line.split('：', 1)[-1].strip()
                    break
        
        # 提取测试方法
        test_methods = re.findall(r'def (test_\w+)\(self.*?\):', script_content)
        for i, method in enumerate(test_methods, 1):
            test_case['test_steps'].append({
                "step": f"执行测试方法 {method}",
                "expected_result": "测试通过，无异常"
            })
        
        # 如果找到了setup和teardown
        if 'def setup_class' in script_content:
            test_case['test_steps'].insert(0, {
                "step": "执行测试环境初始化",
                "expected_result": "环境配置成功"
            })
        
        if 'def teardown_class' in script_content:
            test_case['test_steps'].append({
                "step": "清理测试环境",
                "expected_result": "环境清理成功"
            })
        
        return test_case
    
    def summarize_methods_batch(self, methods: List[Dict]) -> List[Tuple[int, str]]:
        """
        批量生成方法功能总结
        
        Args:
            methods: 方法列表，每个包含method_info
        
        Returns:
            [(index, summary), ...] 返回索引和对应的功能总结
        """
        if not methods:
            return []
        
        print(f"开始批量生成{len(methods)}个方法的功能总结...")
        start_time = time.time()
        
        results = []
        
        # 使用线程池并发处理
        with ThreadPoolExecutor(max_workers=self.max_concurrent_requests) as executor:
            # 提交所有任务
            future_to_index = {}
            
            for i, method in enumerate(methods):
                # 为避免过载，添加小延迟
                if i > 0 and i % self.max_concurrent_requests == 0:
                    time.sleep(self.request_delay)
                
                # 提交任务
                future = executor.submit(
                    self._summarize_method_with_retry,
                    method.get('method_info', {})
                )
                future_to_index[future] = i
            
            # 收集结果
            completed = 0
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    summary = future.result()
                    results.append((index, summary))
                    completed += 1
                    if completed % 10 == 0:
                        print(f"已完成 {completed}/{len(methods)} 个方法总结")
                except Exception as e:
                    print(f"处理方法 {index} 时出错: {e}")
                    # 使用备用总结
                    results.append((index, self._fallback_method_summary(
                        methods[index].get('method_info', {})
                    )))
        
        elapsed_time = time.time() - start_time
        print(f"批量生成完成，耗时: {elapsed_time:.2f}秒")
        
        # 按原始索引排序
        results.sort(key=lambda x: x[0])
        
        return results
    
    def _summarize_method_with_retry(self, method_info: Dict, max_retries: int = 3) -> str:
        """
        带重试的方法功能总结
        """
        for attempt in range(max_retries):
            try:
                prompt = self._build_method_summary_prompt(method_info)
                response = self._call_api(prompt)
                content = self._parse_response(response)
                
                # 返回总结内容
                return content.strip()
                
            except Exception as e:
                if attempt < max_retries - 1:
                    time.sleep(self.request_delay * (attempt + 1))
                else:
                    raise e
        
        return self._fallback_method_summary(method_info)
    
    def _build_method_summary_prompt(self, method_info: Dict) -> str:
        """构建方法功能总结的prompt"""
        method_name = method_info.get('method_name', 'unknown')
        method_type = method_info.get('method_type', 'unknown')
        class_name = method_info.get('class_name', '')
        docstring = method_info.get('docstring', '')
        code = method_info.get('code', '')
        
        # 截断过长的代码
        max_length = 3000
        if len(code) > max_length:
            code = code[:max_length] + "\n... (代码过长，已截断) ..."
        
        context = f"类名：{class_name}" if class_name else "顶层函数"
        
        prompt = f"""请分析以下Python测试方法，生成简洁的功能总结。

方法名：{method_name}
方法类型：{method_type}
{context}

方法代码：
```python
{code}
```

请生成该函数的功能总结，结合使用的主要命令来说明该方法的主要功能和测试目的,并在说明时给出该命令,比如使用...命令...。
如果是test方法，说明它测试什么功能；
如果是setup/teardown方法，说明它准备或清理什么；
如果是辅助方法，说明它提供什么功能。

直接返回总结内容，不要包含其他说明。"""
        
        return prompt
    
    def _fallback_method_summary(self, method_info: Dict) -> str:
        """降级方案：基于规则生成方法总结"""
        method_name = method_info.get('method_name', 'unknown')
        method_type = method_info.get('method_type', 'unknown')
        docstring = method_info.get('docstring', '')
        code = method_info.get('code', '')
        
        # 如果有docstring，使用第一行
        if docstring:
            first_line = docstring.strip().split('\n')[0]
            if first_line:
                return first_line[:50]
        
        # 根据方法类型和名称生成默认总结
        if method_type == 'test':
            # 从方法名提取信息
            parts = method_name.replace('test_', '').split('_')
            return f"测试{' '.join(parts)}功能"
        elif method_type == 'setup':
            return "初始化测试环境和配置"
        elif method_type == 'teardown':
            return "清理测试环境和资源"
        else:
            return f"{method_name}辅助方法"
    
    def _fallback_summary(self, code_snippets: List[str]) -> str:
        """
        降级方案：使用规则生成简单总结
        当LLM不可用时使用
        """
        if not code_snippets:
            return "无代码片段"
        
        # 分析第一个代码片段提取关键信息
        first_snippet = code_snippets[0]
        keywords = []
        steps = []
        
        # 识别关键操作
        if 'CheckCommand' in first_snippet:
            keywords.append('命令检查')
            if 'display' in first_snippet:
                steps.append('执行display命令')
            if 'expect' in first_snippet:
                steps.append('验证预期结果')
        
        if 'send' in first_snippet:
            keywords.append('发送命令')
            if 'system-view' in first_snippet:
                steps.append('进入系统视图')
            if 'interface' in first_snippet:
                steps.append('配置接口')
        
        if 'route' in first_snippet.lower():
            keywords.append('路由')
        
        if 'create_interface' in first_snippet:
            keywords.append('创建接口')
        
        if 'delete_interface' in first_snippet:
            keywords.append('删除接口')
        
        # 构建总结
        if keywords:
            function = f"执行{'/'.join(keywords)}相关操作"
        else:
            function = "执行测试操作"
        
        if steps:
            steps_text = "、".join(steps)
        else:
            steps_text = "执行测试步骤"
        
        # 统计信息
        num_snippets = len(code_snippets)
        
        return f"功能：{function}\n步骤：{steps_text}"
    
    def analyze_pattern(self, pattern: Dict) -> str:
        """
        分析模式并生成详细描述
        
        Args:
            pattern: 包含模式信息的字典
        
        Returns:
            模式的详细描述
        """
        # 提取模式中的代码示例
        examples = pattern.get('examples', [])
        if not examples:
            return "无可分析的代码示例"
        
        # 准备上下文
        context = {
            'frequency': pattern.get('frequency', 0),
            'pattern_type': pattern.get('type', 'unknown'),
            'num_params': len(pattern.get('parameters', [])),
            'files_count': len(set(pattern.get('files', [])))
        }
        
        # 调用总结功能
        summary = self.summarize_code_module(examples[:3], context)
        
        # 添加额外的统计信息
        stats = f"\n\n统计信息：在{context['files_count']}个文件中出现{context['frequency']}次"
        
        return summary + stats
    
    def batch_summarize(self, patterns: List[Dict], max_patterns: int = 10) -> List[Dict]:
        """
        批量总结多个模式
        
        Args:
            patterns: 模式列表
            max_patterns: 最多处理的模式数量
        
        Returns:
            添加了总结的模式列表
        """
        summarized_patterns = []
        
        for i, pattern in enumerate(patterns[:max_patterns]):
            print(f"正在总结模式 {i+1}/{min(len(patterns), max_patterns)}...")
            
            try:
                pattern['llm_summary'] = self.analyze_pattern(pattern)
            except Exception as e:
                print(f"总结模式 {pattern.get('id', i)} 时出错: {e}")
                pattern['llm_summary'] = self._fallback_summary(
                    pattern.get('examples', [])
                )
            
            summarized_patterns.append(pattern)
        
        # 添加未处理的模式（不生成总结）
        if len(patterns) > max_patterns:
            for pattern in patterns[max_patterns:]:
                pattern['llm_summary'] = "未生成总结（超出处理限制）"
                summarized_patterns.append(pattern)
        
        return summarized_patterns